# CNN用于多参数估计
import math
import torch
import torch.nn as nn
import torch.nn.functional as F

class MpeCnnModel(nn.Module):
    def __init__(self):
        super().__init__()
        # 第1层：256*256*8 -> 128*128*16
        # 第2层：128*128*16 -> 64*64*32
        # 第3层：64*64*32 -> 32*32*64
        # 第4层：32*32*64 -> 16*16*128
        # 第5层：16*16*128 -> 8*8*128
        # 第6层：8*8*128 -> 4*4*256
        # 第7层：4*4*256 -> 2*2*512
        # 第8层：2*2*512 -> 2048
        # 第9层：2048 -> 512
        # 第10层：512 -> 12
        self.gen_l1()
        # 第2层：64*64*16 -> 16*16*32
        self.gen_l2()
        # 第3层：16*16*32 -> 4*4*64
        self.gen_l3()
        # 拉直
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(4 * 4 * 64, 256)
        self.fc2 = nn.Linear(256, 12)

    def forward(self, x):
        a1_real = self.l1_mp(self.l1_relu(self.l1_bn(self.l1_conv2d(torch.real(x).float()))))
        a1_imag = self.l1_mp(self.l1_relu(self.l1_bn(self.l1_conv2d(torch.imag(x).float()))))
        a2_real = self.l2_mp(self.l2_relu(self.l2_bn(self.l2_conv2d(a1_real))))
        a2_imag = self.l2_mp(self.l2_relu(self.l2_bn(self.l2_conv2d(a1_imag))))
        a3_real = self.l3_mp(self.l3_relu(self.l3_bn(self.l3_conv2d(a2_real))))
        a3_imag = self.l3_mp(self.l3_relu(self.l3_bn(self.l3_conv2d(a2_imag))))
        a4_real = self.flatten(a3_real)
        a4_imag = self.flatten(a3_imag)
        a5_real = F.relu(self.fc1(a4_real))
        a5_imag = F.relu(self.fc1(a4_imag))
        a6_real = F.relu(self.fc2(a5_real))
        a6_imag = F.relu(self.fc2(a5_imag))
        return a6_real + a6_imag

    def gen_l1(self):
        in_channels = 8
        out_channels = 16
        kernel = (31, 31)
        stride = (1,1)
        dilation = 1
        self.l1_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l1_bn = nn.BatchNorm2d(out_channels)
        self.l1_relu = nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l1_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l2(self):
        in_channels = 16
        out_channels = 32
        kernel = (7, 7)
        stride = (1, 1)
        dilation = 1
        self.l2_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l2_bn = nn.BatchNorm2d(out_channels)
        self.l2_relu = nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l2_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)

    def gen_l3(self):
        in_channels = 32
        out_channels = 64
        kernel = (3, 3)
        stride = (1, 1)
        dilation = 1
        self.l3_conv2d = nn.Conv2d(in_channels,out_channels,kernel,stride=stride, padding='same',dilation=dilation)
        self.l3_bn = nn.BatchNorm2d(out_channels)
        self.l3_relu = nn.ReLU(True)
        mp_size, mp_stride = 4, 4
        self.l3_mp = nn.MaxPool2d(mp_size, mp_stride, padding=0)
